{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"align-center\">\n",
    "<a href=\"https://oumi.ai/\"><img src=\"https://oumi.ai/docs/en/latest/_static/logo/header_logo.png\" height=\"200\"></a>\n",
    "\n",
    "[![Documentation](https://img.shields.io/badge/Documentation-latest-blue.svg)](https://oumi.ai/docs/en/latest/index.html)\n",
    "[![Discord](https://img.shields.io/discord/1286348126797430814?label=Discord)](https://discord.gg/oumi)\n",
    "[![GitHub Repo stars](https://img.shields.io/github/stars/oumi-ai/oumi)](https://github.com/oumi-ai/oumi)\n",
    "</div>\n",
    "\n",
    "👋 Welcome to Open Universal Machine Intelligence (Oumi)!\n",
    "\n",
    "🚀 Oumi is a fully open-source platform that streamlines the entire lifecycle of foundation models - from [data preparation](https://oumi.ai/docs/en/latest/resources/datasets/datasets.html) and [training](hhttps://oumi.ai/docs/en/latest/user_guides/train/train.html) to [evaluation](https://oumi.ai/docs/en/latest/user_guides/evaluate/evaluate.html) and [deployment](https://oumi.ai/docs/en/latest/user_guides/launch/launch.html). Whether you're developing on a laptop, launching large scale experiments on a cluster, or deploying models in production, Oumi provides the tools and workflows you need.\n",
    "\n",
    "🤝 Make sure to join our [Discord community](https://discord.gg/oumi) to get help, share your experiences, and contribute to the project! If you are interested in joining one of the community's open-science efforts, check out our [open collaboration](https://oumi.ai/community) page.\n",
    "\n",
    "⭐ If you like Oumi and you would like to support it, please give it a star on [GitHub](https://github.com/oumi-ai/oumi)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Deploying a Job\n",
    "\n",
    "In this tutorial we'll take a working [JobConfig](https://github.com/oumi-ai/oumi/tree/main/src/oumi/core/configs/job_config.py) and deploy it remotely on a cluster of your choice.\n",
    "\n",
    "This guide dovetails nicely with our [Finetuning Tutorial](https://github.com/oumi-ai/oumi/blob/main/notebooks/Oumi%20-%20Finetuning%20Tutorial.ipynb) where you create your own TrainingConfig and run it locally. Give it a try if you haven't already!\n",
    "\n",
    "We'll cover the following topics:\n",
    "1. Prerequisites\n",
    "1. Choosing a Cloud\n",
    "1. Preparing Your JobConfig\n",
    "1. Launching Your Job\n",
    "1. \\[Advanced\\] Deploying a Training Config"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Oumi Installation\n",
    "First, let's install Oumi. You can find more detailed instructions [here](https://oumi.ai/docs/en/latest/get_started/installation.html). \n",
    "\n",
    "If you have a GPU, you can run the following commands to install Oumi:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install uv -q\n",
    "!uv pip install oumi[gpu] --no-progress --system"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "❗**WARNING:** After the first `pip install`, you may have to restart the notebook for the package updates to take effect (Colab Menu: `Runtime` -> `Restart Session`)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating our working directory\n",
    "For our experiments, we'll use the following folder to save our configs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from pathlib import Path\n",
    "\n",
    "tutorial_dir = \"deploy_training_tutorial\"\n",
    "\n",
    "Path(tutorial_dir).mkdir(parents=True, exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Choosing a Cloud"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll be using the Oumi Launcher to run remote training. To use the launcher, you need to specify which cloud you'd like to run training on.\n",
    "We'll list the clouds below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import oumi.launcher as launcher\n",
    "\n",
    "# Print all available clouds\n",
    "print(launcher.which_clouds())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Local Cloud\n",
    "If you don't have any clouds set up yet, feel free to use the `local` cloud. This will simply execute your job on your current device as if it's a remote cluster. Hardware requirements are ignored for the `local` cloud.\n",
    "\n",
    "#### Other Providers\n",
    "Note that to use a cloud you must already have an account registered with that cloud provider.\n",
    "\n",
    "For example, GCP, RunPod, and Lambda require accounts with billing enabled. Polaris requires an account set up with [ALCF](https://www.alcf.anl.gov/polaris).\n",
    "\n",
    "Once you've picked a cloud, move on to the next step."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparing Your JobConfig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's get started by creating your JobConfig. In the config below, feel free to change `cloud: local` to the cloud you chose in the previous step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile $tutorial_dir/job.yaml\n",
    "\n",
    "name: job-tutorial\n",
    "resources:\n",
    "  cloud: local\n",
    "  # Accelerators is ignored for the local cloud.\n",
    "  accelerators: A100\n",
    "\n",
    "# Upload working directory to remote.\n",
    "# If on the local cloud, we CD into the working directory before running the job.\n",
    "working_dir: .\n",
    "\n",
    "envs:\n",
    "  TEST_ENV_VARIABLE: '\"Hello, World!\"'\n",
    "  OUMI_LOGGING_DIR: \"deploy_training_tutorial/logs\"\n",
    "\n",
    "# `setup` will always be executed before `run`.\n",
    "# No setup is required for this job.\n",
    "#setup: |\n",
    "#  echo \"Running setup...\"\n",
    "\n",
    "run: |\n",
    "  set -e  # Exit if any command failed.\n",
    "\n",
    "  echo \"$TEST_ENV_VARIABLE\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Launching Your Job"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First let's load your JobConfig:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read our JobConfig from the YAML file.\n",
    "job_config = launcher.JobConfig.from_yaml(str(Path(tutorial_dir) / \"job.yaml\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At any point you can easily change the cloud where your job will run by modifying the job's `resources.cloud` parameter:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Manually set the cloud to use.\n",
    "job_config.resources.cloud = \"local\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once you have a job config, kicking off your job is simple:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can optionally specify a cluster name here. If not specified, a random name will\n",
    "# be generated. This is also useful for launching multiple jobs on the same cluster.\n",
    "cluster_name = None\n",
    "\n",
    "# Launch the job!\n",
    "cluster, job_status = launcher.up(job_config, cluster_name)\n",
    "print(f\"Job status: {job_status}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Don't worry if you see any errors from `launcher.up`--you may need to configure permissions to run a job on your specified cloud. The error message should provide you with the proper command to run to authenticate (for GCP this is often `gcloud auth application-default login`).\n",
    "\n",
    "We can quickly check on the status of our job using the `cluster` returned in the previous command:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "while job_status and not job_status.done:\n",
    "    print(\"Job is running...\")\n",
    "    time.sleep(15)\n",
    "    job_status = cluster.get_job(job_status.id)\n",
    "print(\"Job is done!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the job was run on the local cluster, we can view the logs below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logs_dir = Path(tutorial_dir) / \"logs\"\n",
    "for log_file in logs_dir.iterdir():\n",
    "    print(f\"Log file: {log_file}\")\n",
    "    with open(log_file) as f:\n",
    "        print(f.read())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we're done with the cluster, let's turn it down to stop billing for non-local clouds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster.down()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## \\[Advanced\\] Deploying a Training Config"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In our [Finetuning Tutorial](https://github.com/oumi-ai/oumi/blob/main/notebooks/Oumi%20-%20Finetuning%20Tutorial.ipynb), we created and saved a TrainingConfig. We then invoked training by running\n",
    "```shell\n",
    "oumi train -c \"$tutorial_dir/train.yaml\"\n",
    "```\n",
    "\n",
    "You can also run that command as a job! Simply update the \"run\" section of the JobConfig with your desired command:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "path_to_your_train_config = Path(tutorial_dir) / \"train.yaml\"  # Make sure this exists!\n",
    "\n",
    "# Set the `run` command to run your training script.\n",
    "job_config.run = f'oumi train -c \"{path_to_your_train_config}\"'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now your job will run your training config when executed!\n",
    "\n",
    "For a more in-depth overview of the fields in JobConfig, please see our [Running Jobs Remotely tutorial](https://github.com/oumi-ai/oumi/blob/main/notebooks/Oumi%20-%20Running%20Jobs%20Remotely.ipynb)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "oumi",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
